/**
 * @file main.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "linspaced_custom_tiling.h"
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_linspaced_custom.h"
#include "tiling/platform/platform_ascendc.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void linspaced_custom(GM_ADDR startGM, 
                                                       GM_ADDR endGM, 
                                                       const uint32_t n, 
                                                       GM_ADDR zGM,
                                                       LinspacedCustomTilingData tiling);
#endif
extern void GenerateTilingData(uint8_t *tilingBuf, uint32_t blockDim);

int32_t main(int32_t argc, char *argv[])
{
    constexpr uint32_t BLOCK_DIM = 8;
    constexpr uint32_t INPUT_TYPE_SIZE[] = {2, 4, 2, 4};  // half/float/int16/int32
    constexpr uint32_t OUTPUT_TYPE_SIZE = 4;              // 输出固定为float（4字节）
    uint8_t *tiling = nullptr;
    size_t tilingSize =sizeof(LinspacedCustomTilingData);
    
#ifdef ASCENDC_CPU_DEBUG
    tiling = (uint8_t *)AscendC::GmAlloc(tilingSize);
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *startHost, *endHost, *zHost;
    uint8_t *startDevice, *endDevice, *zDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&tiling), tilingSize));
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);
#endif
    GenerateTilingData(tiling, BLOCK_DIM);
    size_t nByteSize = sizeof(uint32_t);
    size_t inputByteSize = INPUT_TYPE_SIZE[reinterpret_cast<LinspacedCustomTilingData *>(tiling)->dataType];
    size_t outputByteSize = reinterpret_cast<LinspacedCustomTilingData *>(tiling)->totalLengthAligned * OUTPUT_TYPE_SIZE;
    size_t validOutputByteSize = reinterpret_cast<LinspacedCustomTilingData *>(tiling)->totalLength * OUTPUT_TYPE_SIZE;
    uint32_t *ns = (uint32_t *)malloc(nByteSize);
    ReadFile("./input/ns.bin", nByteSize, ns, nByteSize); 
    const uint32_t n = ns[0];   
    free(ns);
    
#ifdef ASCENDC_CPU_DEBUG
    constexpr uint32_t DATA_BLOCK_BYTE = 32;
    uint8_t *start = (uint8_t *)AscendC::GmAlloc((inputByteSize + DATA_BLOCK_BYTE - 1) / DATA_BLOCK_BYTE * DATA_BLOCK_BYTE);
    uint8_t *end = (uint8_t *)AscendC::GmAlloc((inputByteSize + DATA_BLOCK_BYTE - 1) / DATA_BLOCK_BYTE * DATA_BLOCK_BYTE);
    uint8_t *z = (uint8_t *)AscendC::GmAlloc((outputByteSize + DATA_BLOCK_BYTE - 1) / DATA_BLOCK_BYTE * DATA_BLOCK_BYTE);

    ReadFile("./input/starts.bin", inputByteSize, start, inputByteSize);
    ReadFile("./input/ends.bin", inputByteSize, end, inputByteSize);
    
    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    
    ICPU_RUN_KF(linspaced_custom, BLOCK_DIM, start, end, n, z, *reinterpret_cast<LinspacedCustomTilingData *>(tiling));

    WriteFile("./output/output_z.bin", z, validOutputByteSize);

    AscendC::GmFree((void *)start);
    AscendC::GmFree((void *)end);
    AscendC::GmFree((void *)z);
    AscendC::GmFree((void *)tiling);
#else
    CHECK_ACL(aclrtMallocHost((void **)(&startHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&endHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&zHost), outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&startDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&endDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&zDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/starts.bin", inputByteSize, startHost, inputByteSize);
    ReadFile("./input/ends.bin", inputByteSize, endHost, inputByteSize);
    
    CHECK_ACL(aclrtMemcpy(startDevice, inputByteSize, startHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(endDevice, inputByteSize, endHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    ACLRT_LAUNCH_KERNEL(linspaced_custom)(BLOCK_DIM, stream, 
                                          startDevice, endDevice, n, zDevice, 
                                          reinterpret_cast<LinspacedCustomTilingData *>(tiling));
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(zHost, outputByteSize, zDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_z.bin", zHost, validOutputByteSize);
    CHECK_ACL(aclrtFree(startDevice));
    CHECK_ACL(aclrtFree(endDevice));
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFreeHost(startHost));
    CHECK_ACL(aclrtFreeHost(endHost));
    CHECK_ACL(aclrtFreeHost(zHost));
    CHECK_ACL(aclrtFreeHost(tiling));
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}
